﻿using System;
using ManagedCuda.BasicTypes;
using OpenTK;
using ManagedCuda;
using ManagedCuda.VectorTypes;
using OpenTK.Graphics;
using OpenTK.Graphics.OpenGL;

namespace ManagedCudaTest
{
    class Program : GameWindow
    {
        static void Main()
        {
            new Program().Run();
        }

        private const float GalaxyMass = 0.000001f;
        private const float ParticleMass = 0.0000000001f;
        private const int GravParticles = 8;
        private readonly int _n;
        private readonly CudaContext _gpu;
        private readonly CudaKernel _kernel;
        private readonly CudaKernel _drawer;
        private readonly uint _glPos;
        private CUgraphicsResource _glPosCuda;
        private readonly CudaDeviceVariable<float3> _dPos;
        private readonly CudaDeviceVariable<float3> _dVel;

        private int _fps;
        private int _lastFpsSecond;

        public Program()
            : base(512, 512, GraphicsMode.Default, "ManagedCudaTest", GameWindowFlags.Default, DisplayDevice.Default, 0, 0, GraphicsContextFlags.ForwardCompatible)
        {
            _n = (int)Math.Pow(2, 15);
            Console.Write("Booting with " + _n + " particles");
            _gpu = CudaContext.CreateOpenGLContext(0, CUCtxFlags.SchedAuto);
            _kernel = _gpu.LoadKernelPTX("kernel.ptx", "kernel");
            _kernel.BlockDimensions = _kernel.MaxThreadsPerBlock;
            _kernel.GridDimensions = _n / _kernel.MaxThreadsPerBlock;
            _drawer = _gpu.LoadKernelPTX("kernel.ptx", "drawer");
            _drawer.BlockDimensions = _drawer.MaxThreadsPerBlock;
            _drawer.GridDimensions = _n / _drawer.MaxThreadsPerBlock;
            Console.WriteLine(", " + (_n / _kernel.MaxThreadsPerBlock) + " blocks, and " + _kernel.MaxThreadsPerBlock + " threads/block");
            Console.WriteLine("Loaded kernel");
            var hPos = new float3[_n];
            var hVel = new float3[_n];
            var rand = new Random();
            for (var i = 0; i < _n; i++)
                GenerateStar(rand, out hPos[i], out hVel[i]);
            Console.WriteLine("Generated star points");
            _dPos = hPos;
            _dVel = hVel;
            Console.WriteLine("Transferred star points to cuda GPU memory");
            GL.ClearColor(0f, 0f, 0f, 1f);
            GL.Color4(1f, 1f, 1f, 1f);
            GL.Viewport(0, 0, Width, Height);
            GL.GenBuffers(1, out _glPos);
            GL.BindBuffer(BufferTarget.ArrayBuffer, _glPos);
            GL.BufferData(BufferTarget.ArrayBuffer, new IntPtr(_n * sizeof(float) * 2), IntPtr.Zero, BufferUsageHint.DynamicDraw);
            GL.Enable(EnableCap.VertexArray);
            Console.WriteLine("Did misc OpenGL stuff");
            _glPosCuda = new CUgraphicsResource();
            HandleError(OpenGLNativeMethods.CUDA3.cuGraphicsGLRegisterBuffer(ref _glPosCuda, _glPos, CUGraphicsRegisterFlags.WriteDiscard));
            Console.WriteLine("Registered OpenGL buffer with CUDA");
            Console.WriteLine("Done initializing");
        }

        private void GenerateStar(Random rand, out float3 pos, out float3 vel)
        {
            do
            {
                pos.x = (float)rand.NextDouble() * 2 - 1;
                pos.y = (float)rand.NextDouble() * 2 - 1;
                pos.z = (float)rand.NextDouble() * 2 - 1;
            } while (pos.Dot(pos) > 1 || pos.x * pos.x + pos.y * pos.y < 0.01);
            pos *= pos.Dot(pos);
            pos.z /= 8;
            pos.z *= pos.z * pos.z;
            var acceleration = (GalaxyMass + ParticleMass * _n / GravParticles * pos.Dot(pos)) / pos.Dot(pos);
            var velMag = (float)Math.Sqrt(pos.Length * acceleration);
            var velDir = float3.Cross(float3.Normalize(pos), new float3(0, 0, 1)); // because I am lazy.
            vel = velDir * velMag;
        }

        protected override void OnUpdateFrame(FrameEventArgs e)
        {
            _kernel.Run(_dPos.DevicePointer, _dVel.DevicePointer, _n);
            base.OnUpdateFrame(e);
        }

        protected override void OnRenderFrame(FrameEventArgs e)
        {
            _fps++;
            if (DateTime.Now.Second != _lastFpsSecond)
            {
                Title = _fps + "fps";
                _fps = 0;
                _lastFpsSecond = DateTime.Now.Second;
            }
            GL.Clear(ClearBufferMask.ColorBufferBit | ClearBufferMask.DepthBufferBit);
            HandleError(DriverAPINativeMethods.GraphicsInterop.cuGraphicsMapResources(1, ref _glPosCuda, new CUstream()));
            var glCuPosPointer = new CUdeviceptr();
            var glCuPosPointerSize = new SizeT();
            HandleError(DriverAPINativeMethods.GraphicsInterop.cuGraphicsResourceGetMappedPointer_v2(ref glCuPosPointer, ref glCuPosPointerSize, _glPosCuda));
            _drawer.Run(_dPos.DevicePointer, glCuPosPointer, _n);
            HandleError(DriverAPINativeMethods.GraphicsInterop.cuGraphicsUnmapResources(1, ref _glPosCuda, new CUstream()));
            GL.BindBuffer(BufferTarget.ArrayBuffer, _glPos);
            GL.VertexPointer(2, VertexPointerType.Float, 0, 0);
            GL.DrawArrays(BeginMode.Points, 0, _n);
            SwapBuffers();
            base.OnRenderFrame(e);
        }

        private static void HandleError(CUResult result)
        {
            if (result != CUResult.Success)
                throw new Exception(result.ToString());
        }
    }
}
